from sqlalchemy.orm import Session

from services.base_prod_common_service import BaseProdCommonService


class ProdService(BaseProdCommonService):
    def __init__(self, model_class, time_field):
        super().__init__(model_class=model_class, time_field=time_field)

    def get_total_ids(self, start_time, end_time, session: Session = None):
        # 如果未提供session，则获取新的session
        if session is None:
            session = self.get_session()
        return (session.query(self.model_class.id)
                .filter(getattr(self.model_class, self.time_field).between(start_time, end_time))
                .all())

    def get_data_by_ids(self, ids, session=None):
        # 如果未提供session，则获取新的session
        if session is None:
            session = self.get_session()
        #     写一个 in 查询
        return (session.query(self.model_class).filter(getattr(self.model_class, 'id').in_(ids))
                .all())
